import numpy as np
import sys
sys.path.append("../") # 将上一级目录也添加到模块的搜索目录中
from common.metrics import r2_score

# 使用向量化运算，提高运算性能
class SimpleLinearRegression2:
    def __init__(self):
        """初始化 线性回归 模型"""
        self.a_ = None
        self.b_ = None

    def fit(self, x_train, y_train):
        """根据训练数据集x_train, y_train训练 线性回归 模型"""
        assert x_train.ndim == 1, "当前线性回归模型只能支持仅有一个特征的训练数据"
        assert len(x_train) == len(y_train), "训练数据个数和结果个数必须相等"

        # x的均值
        x_mean = np.mean(x_train)
        # y的均值
        y_mean = np.mean(y_train)

        # 分子
        num = 0.0
        # 分母
        d = 0.0

        num += (x_train - x_mean).dot(y_train - y_mean)
        d += np.sum((x_train - x_mean) ** 2)

        self.a_ = num / d
        self.b_ = y_mean - self.a_ * x_mean

        return self

    def predict(self, x_predict):
        """给定待预测数据集x_predict，返回x_predict的结果向量"""
        assert x_predict.ndim == 1, "待预测数据集必须为1维向量"
        assert self.a_ is not None and self.b_ is not None, "必须先调用fit方法传入训练数据"

        return self.a_ * x_predict + self.b_

    def score(self, x_test, y_test):
        """根据测试数据集x_test和y_test确定当前模型的准确度"""

        y_predict = self.predict(x_test)
        return r2_score(y_test, y_predict)

    def __repr__(self):
        return "SimpleLinearRegression1()"